#include "parser.h"

#include <algorithm>
#include <iterator>
#include <limits>
#include <list>

#include "ast.h"
#include "lexer.h"
#include "token.h"
#include "types.h"
#include "utils.h"

#define INT_MAX (std::numeric_limits<int>().max())

using BaseToken = Token::BaseToken;
using namespace Utils;
using namespace AST;

Location Parser::Loc() const {
    if (it == tokens.begin()) return {0, 0};
    return (*std::prev(it))->loc;
}

inline BaseToken* Parser::Expect(Tag tag) {
    Tag cur = Peek()->tag;
    if (cur == tag) return Next();
    Err("expected tag %x but got %x", (int)tag, (int)cur);
    return nullptr;
}

Block* Parser::ParseBlock() {
    auto* ret = NewNode<Block>();
    while (!Empty() && !Try((Tag)'}')) {
        if (Check<Token::TypeDecl>()) {
            bool isConst = Try(Tag::ConstQualifier);
            Tag type = Next()->tag;
            if (type != Tag::VoidType && type != Tag::IntType &&
                type != Tag::FloatType) {
                Err("expect a typename");
            }
            auto decls = ParseVarDecl(isConst, type);
            for (auto decl : decls) ret->Add(decl);
        } else
            ret->Add(ParseStmt());
    }
    return ret;
}

ASTNode* Parser::ParseStmt() {
    if (Try((Tag)'{')) {
        return ParseBlock();
    }
    if (Try(Tag::Break)) {
        auto* ret = NewNode<BreakStmt>();
        Expect((Tag)';');
        return ret;
    }
    if (Try(Tag::Continue)) {
        auto* ret = NewNode<ContinueStmt>();
        Expect((Tag)';');
        return ret;
    }
    if (Try(Tag::Return)) {
        auto* ret = NewNode<ReturnStmt>();
        ret->ret = ParseExpr();
        Expect((Tag)';');
        return ret;
    }
    if (Try(Tag::If)) {
        return ParseIfStmt();
    }
    if (Try(Tag::While)) {
        return ParseWhileStmt();
    }

    auto* ret = ParseExpr();
    Expect((Tag)';');
    return ret;
}

IfStmt* Parser::ParseIfStmt() {
    auto* ret = NewNode<IfStmt>();
    Expect((Tag)'(');
    ret->cond = ParseExpr((Tag)')');
    Expect((Tag)')');
    ret->tVal = ParseStmt();
    if (Try(Tag::Else))
        ret->fVal = ParseStmt();
    else
        ret->fVal = NewNode<VoidExpr>();
    return ret;
}

WhileStmt* Parser::ParseWhileStmt() {
    auto* ret = NewNode<WhileStmt>();
    Expect((Tag)'(');
    ret->cond = ParseExpr((Tag)')');
    Expect((Tag)')');
    ret->content = ParseStmt();
    return ret;
}

std::vector<ASTNode*> Parser::ParseDecl() {
    bool isConst = Try(Tag::ConstQualifier);

    Tag type = Next()->tag;
    // Expect(Tag::Identifier);
    if (Next()->tag != Tag::Identifier) {
        Err("expect an identifier");
    }

    bool isFunc = Check((Tag)'(');
    Prev();

    if (isFunc) return {ParseFuncDecl(type)};
    return ParseVarDecl(isConst, type);
}

std::vector<ASTNode*> Parser::ParseVarDecl(bool isConst, Tag baseType) {
    std::vector<ASTNode*> ret;
    Type::Basic irBaseType;
    switch (baseType) {
        case Tag::IntType:
            irBaseType = Type::Basic::Int;
            break;
        case Tag::FloatType:
            irBaseType = Type::Basic::Float;
            break;
        default:
            Err("expect int or float type");
    }

    while (!Try((Tag)';')) {
        if (!Check(Tag::Identifier)) Err("expect an identifier");
        std::string name = Next()->As<Token::Identifier>()->name;
        auto* var = NewNode<VarDecl>(name);
        var->isConst = isConst;
        var->baseType = irBaseType;

        while (Try((Tag)'[')) {
            auto* shape = ParseExpr((Tag)']');
            if (!shape) shape = NewNode<LiteralExpr>(ArithValue(0));
            var->shape.push_back(shape);
            Expect((Tag)']');
        }
        std::reverse(var->shape.begin(), var->shape.end());

        if (Try((Tag)'=')) {
            if (Try((Tag)'{'))
                var->init = ParseInitializer();
            else
                var->init = ParseExpr((Tag)',');
        }
        Try((Tag)',');

        ret.push_back(var);
    }
    return ret;
}

InitializerExpr* Parser::ParseInitializer() {
    auto* ret = NewNode<InitializerExpr>();
    while (!Try((Tag)'}')) {
        if (Try((Tag)'{')) {
            ret->Add(ParseInitializer());
            Try((Tag)',');
        } else {
            ret->Add(ParseExpr((Tag)','));
            Try((Tag)',');
        }
    }
    return ret;
}

FuncDecl* Parser::ParseFuncDecl(Tag baseType) {
    std::string name = Next()->As<Token::Identifier>()->name;
    auto* func = NewNode<FuncDecl>(name);
    switch (baseType) {
        case Tag::IntType:
            func->retType = Type::Basic::Int;
            break;
        case Tag::FloatType:
            func->retType = Type::Basic::Float;
            break;
        case Tag::VoidType:
            func->retType = Type::Basic::Void;
            break;
        default:
            Err("expect a typename");
    }
    Expect((Tag)'(');

    while (!Try((Tag)')')) {
        Tag argType = Next()->tag;
        if (!Check(Tag::Identifier)) Err("expect an identifier");

        std::string argName = Next()->As<Token::Identifier>()->name;
        auto* var = NewNode<VarDecl>(argName);
        switch (argType) {
            case Tag::IntType:
                var->baseType = Type::Basic::Int;
                break;
            case Tag::FloatType:
                var->baseType = Type::Basic::Float;
                break;
            default:
                Err("expect int or float type");
        }
        while (Try((Tag)'[')) {
            auto* shape = ParseExpr((Tag)']');
            var->shape.push_back(shape);
            Expect((Tag)']');
        }
        std::reverse(var->shape.begin(), var->shape.end());
        Try((Tag)',');
        func->args.push_back(var);
    }

    Expect((Tag)'{');
    func->content = ParseBlock();
    return func;
}

ASTNode* Parser::ParseExpr(Tag exitTag, int exitPri) {
    static auto GetPri = [](Tag tag) {
        static std::unordered_map<Tag, int> mp{
            {(Tag)'(', 1},       {(Tag)'[', 1},      {(Tag)'!', 2},
            {(Tag)'*', 3},       {(Tag)'/', 3},      {(Tag)'%', 3},
            {(Tag)'+', 4},       {(Tag)'-', 4},      {Tag::Less, 6},
            {Tag::LessEqual, 6}, {Tag::Greater, 6},  {Tag::GreaterEqual, 6},
            {Tag::Equal, 7},     {Tag::NotEqual, 7}, {Tag::LogicAnd, 11},
            {Tag::LogicOr, 12},  {(Tag)'=', 14},     {(Tag)',', 15},
        };
        if (mp.count(tag)) return mp[tag];
        return INT_MAX;
    };
    static auto GetExitPri = [](Tag tag, bool isUnary = false) {
        if (isUnary) return 2;
        static std::unordered_map<Tag, int> mp{
            {(Tag)'!', 2},          {(Tag)'*', 3},       {(Tag)'/', 3},
            {(Tag)'%', 3},          {(Tag)'+', 4},       {(Tag)'-', 4},
            {Tag::Less, 6},         {Tag::LessEqual, 6}, {Tag::Greater, 6},
            {Tag::GreaterEqual, 6}, {Tag::Equal, 7},     {Tag::NotEqual, 7},
            {Tag::LogicAnd, 11},    {Tag::LogicOr, 12},  {(Tag)'=', 15},
            {(Tag)',', 15},
        };
        if (mp.count(tag)) return mp[tag];
        return INT_MAX;
    };

    ASTNode* ret;
    if (Check<Token::Literal>())
        ret = NewNode<LiteralExpr>(Next());
    else if (Check(Tag::Identifier))
        ret = NewNode<IdentifierExpr>(Next());
    else if (auto* op = Peek()->As<Token::Operator>(); op && op->isUnary()) {
        auto* cur = NewNode<OperatorExpr>(op);
        cur->lhs = NewNode<VoidExpr>();
        cur->rhs = ParseExpr(exitTag, GetExitPri(Next()->tag, true));
        ret = cur;
    } else
        ret = NewNode<VoidExpr>();

    while (!Empty() && !Check(exitTag)) {
        Tag tag = Peek()->tag;
        if (GetPri(tag) >= exitPri) break;

        if (tag == (Tag)'(') {
            Next();
            auto* cur = NewNode<CallExpr>();
            cur->func = ret;
            cur->args = ParseExpr((Tag)')', GetExitPri(tag));
            Expect((Tag)')');
            ret = cur;
        } else if (tag == (Tag)'[') {
            Next();
            auto* cur = NewNode<ArrayExpr>();
            cur->lhs = ret;
            cur->rhs = ParseExpr((Tag)']', GetExitPri(tag));
            Expect((Tag)']');
            ret = cur;
        } else if (tag == (Tag)'=') {
            Next();
            auto* cur = NewNode<AssignExpr>();
            cur->lhs = ret;
            cur->rhs = ParseExpr(exitTag, GetExitPri(tag));
            ret = cur;
        } else {
            auto* cur = NewNode<OperatorExpr>(Next());
            cur->lhs = ret;
            cur->rhs = ParseExpr(exitTag, GetExitPri(tag));
            ret = cur;
        }
    }
    return ret;
}

Block* Parser::Parse() {
    auto* ret = NewNode<Block>();
    while (!Empty() && !Try((Tag)'}')) {
        if (Check<Token::TypeDecl>()) {
            auto decls = ParseDecl();
            for (auto decl : decls) ret->Add(decl);
        } else
            ret->Add(ParseStmt());
    }
    return ret;
};

void* Parser::Run(void* input) {
    tokens = *(std::list<BaseToken*>*)input;
    it = tokens.begin();
    auto* ret = Parse();
    return ret;
}
